import { create } from 'zustand';
import { persist } from 'zustand/middleware';
import { ChatMessage, ChatSession, Document, User } from '@/types';

interface AppState {
  // 用户状态
  user: User | null;
  setUser: (user: User | null) => void;

  // 聊天状态
  currentSession: ChatSession | null;
  sessions: ChatSession[];
  messages: ChatMessage[];
  isLoading: boolean;
  isStreaming: boolean;

  // 聊天操作
  setCurrentSession: (session: ChatSession | null) => void;
  addSession: (session: ChatSession) => void;
  updateSession: (sessionId: string, updates: Partial<ChatSession>) => void;
  deleteSession: (sessionId: string) => void;
  setSessions: (sessions: ChatSession[]) => void;

  addMessage: (message: ChatMessage) => void;
  updateMessage: (messageId: string, updates: Partial<ChatMessage>) => void;
  setMessages: (messages: ChatMessage[]) => void;
  clearMessages: () => void;
  setIsLoading: (loading: boolean) => void;
  setIsStreaming: (streaming: boolean) => void;

  // 文档操作
  documents: Document[];
  selectedDocument: Document | null;
  uploadProgress: Record<string, number>;
  isUploading: boolean;

  setDocuments: (documents: Document[]) => void;
  addDocument: (document: Document) => void;
  updateDocument: (documentId: string, updates: Partial<Document>) => void;
  deleteDocument: (documentId: string) => void;
  setSelectedDocument: (document: Document | null) => void;

  setUploadProgress: (progress: Record<string, number>) => void;
  setIsUploading: (uploading: boolean) => void;

  // UI操作
  sidebarOpen: boolean;
  theme: 'light' | 'dark';

  setSidebarOpen: (open: boolean) => void;
  toggleSidebar: () => void;
  setTheme: (theme: 'light' | 'dark') => void;
  toggleTheme: () => void;

  // 重置状态
  reset: () => void;
}

const initialState = {
  user: null,
  currentSession: null,
  sessions: [],
  messages: [],
  isLoading: false,
  isStreaming: false,
  documents: [],
  selectedDocument: null,
  uploadProgress: {},
  isUploading: false,
  sidebarOpen: true,
  theme: 'light' as const,
};

export const useAppStore = create<AppState>()(
  persist(
    (set, _get) => ({
      ...initialState,

      // 用户操作
      setUser: user => set({ user }),

      // 聊天操作
      setCurrentSession: session => set({ currentSession: session }),

      addSession: session =>
        set(state => ({
          sessions: [...state.sessions, session],
        })),

      updateSession: (sessionId, updates) =>
        set(state => ({
          sessions: state.sessions.map(session =>
            session.id === sessionId ? { ...session, ...updates } : session
          ),
          currentSession:
            state.currentSession?.id === sessionId
              ? { ...state.currentSession, ...updates }
              : state.currentSession,
        })),

      deleteSession: sessionId =>
        set(state => ({
          sessions: state.sessions.filter(session => session.id !== sessionId),
          currentSession:
            state.currentSession?.id === sessionId
              ? null
              : state.currentSession,
        })),

      setSessions: sessions => set({ sessions }),

      addMessage: message =>
        set(state => ({
          messages: [...state.messages, message],
        })),

      updateMessage: (messageId, updates) =>
        set(state => ({
          messages: state.messages.map(msg =>
            msg.id === messageId ? { ...msg, ...updates } : msg
          ),
        })),

      setMessages: messages => set({ messages }),

      clearMessages: () => set({ messages: [] }),

      setIsLoading: isLoading => set({ isLoading }),

      setIsStreaming: isStreaming => set({ isStreaming }),

      // 文档操作
      setDocuments: documents => set({ documents }),

      addDocument: document =>
        set(state => ({
          documents: [...state.documents, document],
        })),

      updateDocument: (documentId, updates) =>
        set(state => ({
          documents: state.documents.map(doc =>
            doc.id === documentId ? { ...doc, ...updates } : doc
          ),
          selectedDocument:
            state.selectedDocument?.id === documentId
              ? { ...state.selectedDocument, ...updates }
              : state.selectedDocument,
        })),

      deleteDocument: documentId =>
        set(state => ({
          documents: state.documents.filter(doc => doc.id !== documentId),
          selectedDocument:
            state.selectedDocument?.id === documentId
              ? null
              : state.selectedDocument,
        })),

      setSelectedDocument: document => set({ selectedDocument: document }),

      setUploadProgress: uploadProgress =>
        set(state => ({
          uploadProgress: { ...state.uploadProgress, ...uploadProgress },
        })),

      setIsUploading: isUploading => set({ isUploading }),

      // UI操作
      setSidebarOpen: sidebarOpen => set({ sidebarOpen }),

      toggleSidebar: () => set(state => ({ sidebarOpen: !state.sidebarOpen })),

      setTheme: theme => set({ theme }),

      toggleTheme: () =>
        set(state => ({
          theme: state.theme === 'light' ? 'dark' : 'light',
        })),

      // 重置状态
      reset: () => set(initialState),
    }),
    {
      name: 'rag-app-storage',
      partialize: state => ({
        user: state.user,
        theme: state.theme,
        sidebarOpen: state.sidebarOpen,
      }),
    }
  )
);

// 选择器函数
export const useUser = () => useAppStore(state => state.user);
export const useCurrentSession = () =>
  useAppStore(state => state.currentSession);
export const useSessions = () => useAppStore(state => state.sessions);
export const useMessages = () => useAppStore(state => state.messages);
export const useDocuments = () => useAppStore(state => state.documents);
export const useSelectedDocument = () =>
  useAppStore(state => state.selectedDocument);
export const useIsLoading = () => useAppStore(state => state.isLoading);
export const useIsStreaming = () => useAppStore(state => state.isStreaming);
export const useIsUploading = () => useAppStore(state => state.isUploading);
export const useUploadProgress = () =>
  useAppStore(state => state.uploadProgress);
export const useSidebarOpen = () => useAppStore(state => state.sidebarOpen);
export const useTheme = () => useAppStore(state => state.theme);

// 操作函数 hooks
export const useChatActions = () => ({
  setCurrentSession: useAppStore(state => state.setCurrentSession),
  addSession: useAppStore(state => state.addSession),
  updateSession: useAppStore(state => state.updateSession),
  deleteSession: useAppStore(state => state.deleteSession),
  setSessions: useAppStore(state => state.setSessions),
  addMessage: useAppStore(state => state.addMessage),
  updateMessage: useAppStore(state => state.updateMessage),
  setMessages: useAppStore(state => state.setMessages),
  clearMessages: useAppStore(state => state.clearMessages),
  setIsLoading: useAppStore(state => state.setIsLoading),
  setIsStreaming: useAppStore(state => state.setIsStreaming),
});

export const useDocumentActions = () => ({
  setDocuments: useAppStore(state => state.setDocuments),
  addDocument: useAppStore(state => state.addDocument),
  updateDocument: useAppStore(state => state.updateDocument),
  deleteDocument: useAppStore(state => state.deleteDocument),
  setSelectedDocument: useAppStore(state => state.setSelectedDocument),
  setUploadProgress: useAppStore(state => state.setUploadProgress),
  setIsUploading: useAppStore(state => state.setIsUploading),
});

export const useUIActions = () => ({
  setSidebarOpen: useAppStore(state => state.setSidebarOpen),
  toggleSidebar: useAppStore(state => state.toggleSidebar),
  setTheme: useAppStore(state => state.setTheme),
  toggleTheme: useAppStore(state => state.toggleTheme),
});
